import csv
import numpy as np
import sys 
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
project_path = os.path.join(current_dir, '..')
sys.path.append(project_path)
sys.path.append(current_dir)
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from scipy.signal import butter, filtfilt, hilbert
import matplotlib.pyplot as plt
from matplotlib import rcParams
from scipy.signal import get_window


def gen_single_freq_wave_from_timescale(x, freq_base):

    # 生成标准正弦波
    # 计算角频率
    omega = 2 * np.pi * freq_base
    # 生成正弦波数据
    wave = np.sin(omega * x)
    return wave

def dc2ac(signal):
    y = signal
    fft_signal = np.fft.fft(y)
    fft_signal[0] = 0; #dc=>ac
    fft_signal_no_phase_correction = np.fft.ifft(fft_signal)
    mofz = np.real(fft_signal_no_phase_correction) #python与matlab不同，实部就是模。虚部是幅角。
    #print("signal(dc)", signal)
    return mofz;

def band_filter(x, signal, freq_low, freq_high):
    saps=1.0/x[1]
    # 定义一个二阶巴特沃斯滤波器
    b, a = butter(2, [2*freq_low/saps, 2*freq_high/saps], btype='band')
    # 应用滤波器
    filtered_signal = filtfilt(b, a, signal)
    signal_abs = [np.abs(z.real) for z in filtered_signal]
    #print("signal(ac)", signal);
    #print("filter_signal", signal);
    return signal_abs

def envelope(signal):
    complex_signal = hilbert(signal);
    signal_abs = [np.abs(z.real) for z in complex_signal]
    return signal_abs;

def lowFilterSignal(x_axis, sig, shutdown_freq):
    #calc_fft
    x = x_axis
    y = sig
    fft_signal = np.fft.fft(y)

    n = np.arange(len(fft_signal)//2)
    freq = np.arange(len(fft_signal)//2)
    freq = freq *(1.0/(len(fft_signal)*x[1]));

    for k in n:
        if freq[k]> shutdown_freq:
           break
    if(k<max(n)):
        fft_signal[k:int(k+2*(max(n)-k)+1)] = 0
        signal_filted = np.fft.ifft(fft_signal)
        signal_abs = [np.abs(z.real) for z in signal_filted]
    else:
        signal_abs= sig
    return (x, signal_abs);


def GetFFTOfSignal(x_axis, sig, window_type='balckman'):
    # 计算采样间隔
    sampling_interval = x_axis[1] - x_axis[0]
    # 计算采样频率
    sampling_frequency = 1.0 / sampling_interval
    # 信号长度
    signal_length = len(sig)
    #calc_fft
    x = x_axis
    y = sig

    # 选择窗函数
    if window_type == 'hanning':
        window = np.hanning(signal_length)
    elif window_type == 'hamming':
        window = np.hamming(signal_length)
    elif window_type == 'blackman':
        window = np.blackman(signal_length)
    else:
        window = None
        raise ValueError("Unsupported window type. Supported types are 'hanning', 'hamming', and 'blackman'.")

    # 加窗
    if window is not None:
        windowed_signal = sig * window
    fft_signal = 2/signal_length*np.fft.fft(windowed_signal)
    freq = np.fft.fftfreq(signal_length, sampling_interval)
    freq = freq[:signal_length//2]
    fft_signal_abs = fft_signal[0:len(fft_signal)//2]
    
    fft_x = freq;
    fft_y = np.abs(fft_signal_abs) #ratio = 2/len(sig)， python complex is [rho, theta]
    fft_y[0] = 0; #try to remove dc
    ffty_toshow = fft_y
    ffty_raw_ac = fft_signal
    ffty_raw_ac[0] = 0  #dc=>ac
    return (fft_x, ffty_toshow, ffty_raw_ac);

def GetFFTOfSignal_withShutdownFreq(x_axis, sig, freq_shutdown, window_type='balckman'):
    #calc_fft
    signal_length=len(sig)

    # 选择窗函数
    if window_type == 'hanning':
        window = np.hanning(signal_length)
    elif window_type == 'hamming':
        window = np.hamming(signal_length)
    elif window_type == 'blackman':
        window = np.blackman(signal_length)
    else:
        window = None
        raise ValueError("Unsupported window type. Supported types are 'hanning', 'hamming', and 'blackman'.")

    # 加窗
    if window is not None:
        windowed_signal = sig * window

    fft_signal = 2/signal_length*np.fft.fft(windowed_signal)
    fft_signal[0] = 0; #dc=>ac
    freq = np.arange(len(fft_signal)//2)
    freq = freq *(1.0/(len(fft_signal)*x_axis[1]));
    print(freq[-1], len(freq))
    fft_signal_abs = fft_signal[0:len(fft_signal)//2]
    
    max_freq_index =next((i for i, num in enumerate(freq) if num > freq_shutdown), len(freq))
    print("max_freq_index = ", max_freq_index, ",values=", freq[max_freq_index])
    fft_x = freq[:max_freq_index];
    fft_y = np.abs(fft_signal_abs[:max_freq_index]);

    ffty_toshow = fft_y
    ffty_raw_ac = fft_signal
    return (fft_x, ffty_toshow, ffty_raw_ac);

def DrawSignals(signals, freq_shutdown, signal_memo):
    cnt = len(signals)
    # 绘制频谱图
    plt.figure(figsize=(12, 6))
    row = cnt;
    col = 2;
    sn = 1;
    for item in signals:
        x = item[0];
        y = item[1];
        sTitle = item[2];
        sX = item[3];
        sY = item[4];

        # 绘制时间域波形
        plt.subplot(row, 2, sn)
        sn = sn +1;
        plt.plot(x, y, color="green")
        plt.xlabel(sX)
        plt.ylabel(sY)
        plt.title(sTitle)

        #calc_fft
        (fft_x, fft_y, dumb) = GetFFTOfSignal_withShutdownFreq(x, y, freq_shutdown)

        #绘制频谱
        plt.subplot(row, 2, sn)
        sn = sn +1
        plt.plot(fft_x, fft_y, color='blue')
        str_fft = 'Freq Spectrum of Left Figure(datafile=%s)' %(signal_memo)
        plt.title(str_fft)
        plt.xlabel('Frequency[cut to %5.1fHz], totalPt=%d'  %(freq_shutdown, len(fft_y)))
        plt.ylabel('Amplitude')
        plt.axis('auto')  # 启用自动调整大小和缩放功能
        plt.grid(True)
    plt.tight_layout()
    #plt.show()
    plt.savefig("jijy.png")

# 自定义坐标轴刻度显示函数
def formater_u16_to_50g(y, pos):
    return '{:.0f}'.format((y - 32768)*50/ 32768)

#used for abs_values, like FFT, envelope.
def formater_u16_to_50g_abs(y, pos):
    return '{:.0f}'.format((y-0)*50/32768)

def DrawU16Signals(signals, freq_shutdown, signal_memo, maxScaleInGravity):
    # 设置中文字体
    rcParams['font.sans-serif'] = ['SimHei']  # 指定中文字体
    rcParams['axes.unicode_minus'] = False  # 解决负号问题

    cnt = len(signals)
    # 绘制频谱图
    plt.figure(figsize=(12, 6))
    row = cnt;
    col = 2;
    sn = 1;
    for item in signals:
        x = item[0];
        y = item[1];
        sTitle = item[2];
        sX = item[3];
        sY = item[4];

        # 绘制时间域波形
        fig1 = plt.subplot(row, 2, sn)
        sn = sn +1;
        fig1.plot(x, y, color="green")
        fig1.set_xlabel(sX)
        fig1.set_ylabel(sY)
        fig1.set_title(sTitle)
        
        # 设置Y轴显示
        formatter = ticker.FuncFormatter(formater_u16_to_50g)
        fig1.yaxis.set_major_formatter(formatter)
        fig1.yaxis.set_major_locator(ticker.AutoLocator())

        #calc_fft
        (fft_x, fft_y, raw_fft) = GetFFTOfSignal_withShutdownFreq(x, y, freq_shutdown)
        # 使用列表推导式对序列中的每个元素乘以系数
        fft_y = [x * maxScaleInGravity/32768/32768 for x in fft_y] #u16=>g

        #绘制频谱
        fig2 = plt.subplot(row, 2, sn)
        sn = sn +1
        fig2.plot(fft_x, fft_y, color='blue')
        str_fft = 'Freq Spectrum of Left Figure(datafile=%s)' %(signal_memo)
        fig2.set_title(str_fft)
        fig2.set_xlabel('Frequency[cut to %5.1fHz], totalPt=%d'  %(freq_shutdown, len(raw_fft)//2))
        fig2.set_ylabel('g(acc)')
        fig2.axis('auto')  # 启用自动调整大小和缩放功能
        fig2.grid(True)
        # 设置Y轴显示
    plt.tight_layout()
    plt.show()
